//! RSA加密实现
//! 
//! 提供RSA加密解密功能

use rsa::{Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey};
use rsa::pkcs1::{EncodeRsaPrivateKey, EncodeRsaPublicKey, DecodeRsaPrivateKey, DecodeRsaPublicKey};
use rand::rngs::OsRng;
use base64::{Engine as _, engine::general_purpose};

/// RSA错误类型
#[derive(Debug, thiserror::Error)]
pub enum RSAError {
    #[error("RSA加密失败: {0}")]
    EncryptionFailed(String),
    #[error("RSA解密失败: {0}")]
    DecryptionFailed(String),
    #[error("RSA密钥生成失败: {0}")]
    KeyGenerationFailed(String),
    #[error("RSA密钥解析失败: {0}")]
    KeyParsingFailed(String),
    #[error("Base64编码失败: {0}")]
    Base64Error(#[from] base64::DecodeError),
}

pub type RSAResult<T> = Result<T, RSAError>;

/// 生成RSA密钥对
/// 
/// # Examples
/// ```
/// use librarys::crypto::asymmetric::generate_rsa_keypair;
/// 
/// let (public_key, private_key) = generate_rsa_keypair().unwrap();
/// assert!(!public_key.is_empty());
/// assert!(!private_key.is_empty());
/// ```
pub fn generate_rsa_keypair() -> RSAResult<(String, String)> {
    let mut rng = OsRng;
    let bits = 2048;
    
    let private_key = RsaPrivateKey::new(&mut rng, bits)
        .map_err(|e| RSAError::KeyGenerationFailed(e.to_string()))?;
    
    let public_key = RsaPublicKey::from(&private_key);
    
    // 将密钥转换为PEM格式
    let private_pem = private_key.to_pkcs1_pem(rsa::pkcs1::LineEnding::LF)
        .map_err(|e| RSAError::KeyGenerationFailed(e.to_string()))?;
    
    let public_pem = public_key.to_pkcs1_pem(rsa::pkcs1::LineEnding::LF)
        .map_err(|e| RSAError::KeyGenerationFailed(e.to_string()))?;
    
    Ok((public_pem.to_string(), private_pem.to_string()))
}

/// RSA加密
/// 
/// # Examples
/// ```
/// use librarys::crypto::asymmetric::{generate_rsa_keypair, rsa_encrypt};
/// 
/// let (public_key, _) = generate_rsa_keypair().unwrap();
/// let plaintext = "hello world";
/// let encrypted = rsa_encrypt(plaintext, &public_key).unwrap();
/// assert!(!encrypted.is_empty());
/// ```
pub fn rsa_encrypt(plaintext: &str, public_key: &str) -> RSAResult<String> {
    let public_key = RsaPublicKey::from_pkcs1_pem(public_key)
        .map_err(|e| RSAError::KeyParsingFailed(e.to_string()))?;
    
    let mut rng = OsRng;
    let data = plaintext.as_bytes();
    
    let encrypted_data = public_key.encrypt(&mut rng, Pkcs1v15Encrypt, data)
        .map_err(|e| RSAError::EncryptionFailed(e.to_string()))?;
    
    // 使用Base64编码返回
    Ok(general_purpose::STANDARD.encode(encrypted_data))
}

/// RSA解密
/// 
/// # Examples
/// ```
/// use librarys::crypto::asymmetric::{generate_rsa_keypair, rsa_encrypt, rsa_decrypt};
/// 
/// let (public_key, private_key) = generate_rsa_keypair().unwrap();
/// let plaintext = "hello world";
/// let encrypted = rsa_encrypt(plaintext, &public_key).unwrap();
/// let decrypted = rsa_decrypt(&encrypted, &private_key).unwrap();
/// assert_eq!(decrypted, plaintext);
/// ```
pub fn rsa_decrypt(ciphertext: &str, private_key: &str) -> RSAResult<String> {
    let private_key = RsaPrivateKey::from_pkcs1_pem(private_key)
        .map_err(|e| RSAError::KeyParsingFailed(e.to_string()))?;
    
    // 解码Base64
    let encrypted_data = general_purpose::STANDARD.decode(ciphertext)?;
    
    let decrypted_data = private_key.decrypt(Pkcs1v15Encrypt, &encrypted_data)
        .map_err(|e| RSAError::DecryptionFailed(e.to_string()))?;
    
    String::from_utf8(decrypted_data)
        .map_err(|e| RSAError::DecryptionFailed(format!("Invalid UTF-8: {}", e)))
}

/// RSA签名（使用SHA256）
/// 
/// # Examples
/// ```
/// use librarys::crypto::asymmetric::{generate_rsa_keypair, rsa_sign};
/// 
/// let (_, private_key) = generate_rsa_keypair().unwrap();
/// let message = "hello world";
/// let signature = rsa_sign(message, &private_key).unwrap();
/// assert!(!signature.is_empty());
/// ```
pub fn rsa_sign(message: &str, private_key: &str) -> RSAResult<String> {
    use rsa::Pkcs1v15Sign;
    use sha2::{Digest, Sha256};
    
    let private_key = RsaPrivateKey::from_pkcs1_pem(private_key)
        .map_err(|e| RSAError::KeyParsingFailed(e.to_string()))?;
    
    let data = message.as_bytes();
    
    // 手动计算消息的哈希值
    let mut hasher = Sha256::new();
    hasher.update(data);
    let hashed = hasher.finalize();
    
    let signature = private_key.sign(Pkcs1v15Sign::new_unprefixed(), &hashed)
        .map_err(|e| RSAError::EncryptionFailed(format!("Sign failed: {}", e)))?;
    
    Ok(general_purpose::STANDARD.encode(signature))
}

/// RSA验签（使用SHA256）
/// 
/// # Examples
/// ```
/// use librarys::crypto::asymmetric::{generate_rsa_keypair, rsa_sign, rsa_verify};
/// 
/// let (public_key, private_key) = generate_rsa_keypair().unwrap();
/// let message = "hello world";
/// let signature = rsa_sign(message, &private_key).unwrap();
/// let is_valid = rsa_verify(message, &signature, &public_key).unwrap();
/// assert!(is_valid);
/// ```
pub fn rsa_verify(message: &str, signature: &str, public_key: &str) -> RSAResult<bool> {
    use rsa::Pkcs1v15Sign;
    use sha2::{Digest, Sha256};
    
    let public_key = RsaPublicKey::from_pkcs1_pem(public_key)
        .map_err(|e| RSAError::KeyParsingFailed(e.to_string()))?;
    
    let signature_bytes = general_purpose::STANDARD.decode(signature)?;
    let data = message.as_bytes();
    
    // 手动计算消息的哈希值
    let mut hasher = Sha256::new();
    hasher.update(data);
    let hashed = hasher.finalize();
    
    match public_key.verify(Pkcs1v15Sign::new_unprefixed(), &hashed, &signature_bytes) {
        Ok(_) => Ok(true),
        Err(_) => Ok(false),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rsa_keypair_generation() {
        let result = generate_rsa_keypair();
        assert!(result.is_ok());
        
        let (public_key, private_key) = result.unwrap();
        assert!(!public_key.is_empty());
        assert!(!private_key.is_empty());
        assert!(public_key.contains("BEGIN RSA PUBLIC KEY"));
        assert!(private_key.contains("BEGIN RSA PRIVATE KEY"));
    }

    #[test]
    fn test_rsa_encrypt_decrypt() {
        let (public_key, private_key) = generate_rsa_keypair().unwrap();
        let plaintext = "hello world";
        
        let encrypted = rsa_encrypt(plaintext, &public_key).unwrap();
        assert!(!encrypted.is_empty());
        
        let decrypted = rsa_decrypt(&encrypted, &private_key).unwrap();
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_rsa_sign_verify() {
        let (public_key, private_key) = generate_rsa_keypair().unwrap();
        let message = "hello world";
        
        let signature = rsa_sign(message, &private_key).unwrap();
        assert!(!signature.is_empty());
        
        let is_valid = rsa_verify(message, &signature, &public_key).unwrap();
        assert!(is_valid);
        
        // 测试错误的签名
        let is_invalid = rsa_verify("different message", &signature, &public_key).unwrap();
        assert!(!is_invalid);
    }

    #[test]
    fn test_rsa_error_cases() {
        // 测试错误的密钥
        let result = rsa_encrypt("test", "invalid key");
        assert!(result.is_err());
        
        let result = rsa_decrypt("test", "invalid key");
        assert!(result.is_err());
        
        // 测试错误的Base64
        let (_, private_key) = generate_rsa_keypair().unwrap();
        let result = rsa_decrypt("invalid base64!", &private_key);
        assert!(result.is_err());
    }
}